import pandas as pd
import numpy as np
from DataControl.orm.pgsql import pgsql
import matplotlib.dates as mdates
from matplotlib.dates import DateFormatter
import matplotlib.pyplot as plt
from mplfinance.original_flavor import candlestick_ohlc
import mplfinance as mpf
from datetime import datetime
from pandas import DateOffset
from Index.utils import shapeOfLine

#获取MACD的指数
class MACD(object):
    """
        MACD 类用于计算和分析股票数据中的 MACD 指标及相关技术指标。该类提供了计算 MACD、RSI 等指标的方法，并能够检测不同类型的背离信号，如波谷底背离、死叉底背离等。

        属性：
            df (DataFrame): 包含股票数据的 Pandas DataFrame，需包含 `close`、`trade_date` 等列。

        方法：
            __init__(df): 初始化 MACD 对象，计算 MACD、RSI 以及其他技术指标。
            MACD(short_window=12, long_window=26, signal_window=9): 计算 MACD 指标及其信号线和柱状图。
            RSI(window=14): 计算相对强弱指标 (RSI)。
            detect_macd_trough(MACD, DEA, target): 检测某个点是否是 MACD 波谷底背离。
            detect_all_macd_trough_valley(): 检测所有的 MACD 波谷底背离点。
            detect_macd_cross_down(MACD, DEA, close, target): 检测某个点是否是 MACD 死叉底背离。
            detect_all_macd_cross_down(): 检测所有的 MACD 死叉底背离点。
            detect_macd_column_down(hist, close, target, span): 检测某个点是否是 MACD 柱线底背离。
            detect_all_macd_column_down(span): 检测所有的 MACD 柱线底背离点。
            detect_macd_top_back(macd, ma10): 检测某个点是否是 MACD 波峰顶背离。
            detect_all_macd_top_back(span): 检测所有的 MACD 波峰顶背离点。
            detect_macd_top_cross(macd, close, dea, target, span): 检测某个点是否是 MACD 高位死叉背离。
            detect_all_macd_top_cross(span): 检测所有的 MACD 高位死叉背离点。
            detect_macd_top_column(hist, close, target, spanOfclose=5, spanOfhist=3): 检测某个点是否是 MACD 柱线顶背离。
            detect_all_macd_top_column(spanOfclose=5, spanOfhist=3): 检测所有的 MACD 柱线顶背离点。
            plot_all_macd_trough(typ): 绘制包含指定类型 MACD 背离信号的图表。
        """

    def __init__(self, df):
        self.df = df
        self.RSI()
        self.MACD()
        self.df['MA5'] = df['close'].rolling(5).mean()
        self.df['MA10'] = df['close'].rolling(5).mean()
        self.df['MOM'] = self.df['close'].diff(periods=10)
        self.df['EMA12'] = self.df['close'].ewm(span=12, adjust=False).mean()
        self.df['EMA26'] = self.df['close'].ewm(span=26, adjust=False).mean()
        #print(self.df)
        self.df = self.df.dropna()
        db = pgsql()


    def MACD(self, short_window=12, long_window=26, signal_window=9):
        #data =

        data = self.df['close']
        # 计算短期EMA
        EMA_short = data.ewm(span=short_window, adjust=False).mean()

        # 计算长期EMA
        EMA_long = data.ewm(span=long_window, adjust=False).mean()

        # 计算MACD线
        self.df['MACD'] = EMA_short - EMA_long

        # 计算信号线
        self.df['DEA'] = self.df['MACD'].ewm(span=signal_window, adjust=False).mean()

        # 计算柱状图
        self.df['Histogram'] = self.df['MACD'] - self.df['DEA']
        #return self.df['MACD'].values

    def RSI(self, window=14):
        data = self.df['close']
        # 计算价格变化
        Change = data.diff()

        # 计算上涨和下跌
        Gain = Change.apply(lambda x: x if x > 0 else 0)
        Loss = Change.apply(lambda x: -x if x < 0 else 0)

        # 计算平均上涨和平均下跌
        Avg_Gain = Gain.rolling(window=window, min_periods=1).mean()
        Avg_Loss = Loss.rolling(window=window, min_periods=1).mean()

        # 计算相对强度RS
        RS = Avg_Gain / Avg_Loss

        # 计算RSI
        self.df['RSI'] = 100 - (100 / (1 + RS))

       # return self.df


    # 检测某个点是否波谷底背离
    def detect_macd_trough(self, MACD, DEA, target):
        """
        DEA: DEA序列
        MACD:  MACD线序列/DIF
        target: 点下标
        当MACD到达谷底后DEA仍然在下降
        此时MACD在上升，而DEA正在下降，方向相反
        """
        flag = True
        flag &= shapeOfLine.isValley(arr = MACD, target = target, span = 3)
        flag &= shapeOfLine.isMonotony_infuture(arr = DEA, target = target, ascent = -1, span = 2)
        return flag

    # df所有的节点是波谷底背离的点
    def detect_all_macd_trough_valley(self):
        MACD = self.df['MACD'].to_numpy()
        DEA = self.df['DEA'].to_numpy()
        date = self.df['trade_date'].to_numpy()
        res = []
        for i in range(len(MACD)) :
            if(self.detect_macd_trough(MACD, DEA, i)) :
                res.append(date[i])
        return res

        # 检测某个点是否是死叉底背离
    def detect_macd_cross_down(self, MACD, DEA, close, target):
        """
        close: 价格曲线
        DEA: DEA序列
        MACD:  MACD线序列/DIF
        target: 点下标
        出现死叉，即DIF线从上向下穿过DEA线
        当死叉后DIF线向下运行，而股价却向上运行
        在0轴上的死叉更为靠谱
        """
        flag = True
        flag &= shapeOfLine.isDeadCross(MACD, DEA, target)
        flag &= shapeOfLine.isMonotony_infuture(close, target, 1, 2)
        flag &= shapeOfLine.isMonotony_infuture(MACD, target, -1, 2)
        for i in range(max(0, target - 3), min(target + 2, len(MACD)), 1):
            flag &= (MACD[i] > 0 and DEA[i] > 0)
        return flag

    # df所有的节点是死叉底背离的点
    def detect_all_macd_cross_down(self):
        MACD = self.df['MACD'].to_numpy()
        DEA = self.df['DEA'].to_numpy()
        date = self.df['trade_date'].to_numpy()
        close = self.df['MA10'].to_numpy()
        res = []
        for i in range(len(MACD)):
            if(self.detect_macd_cross_down(MACD, DEA, close, i)):
                res.append(date[i])
        return res;

    def detect_macd_column_down(self, hist, close, target, span):
        """
        close: 价格曲线
        hist: MACD柱
        target: 点下标
        MACD绿柱增长，股价上涨；
        MACD红柱增长，股价下跌；
        span: 判断过去多久的点
        """
        isGreen = True
        isRed = True
        for i in range(max(0, target - span), target + 1, 1):
            if hist[i] > 0: isGreen = False
            if hist[i] < 0: isRed = False
        flag = isGreen | isRed
        if isGreen:
            flag &= shapeOfLine.isMonotony_inpast(hist, target, 1, span)
            flag &= shapeOfLine.isMonotony_inpast(close, target, 1, span)
        if isRed:
            flag &= shapeOfLine.isMonotony_inpast(hist, target, 1, span)
            flag &= shapeOfLine.isMonotony_inpast(close, target, -1, span)
        return flag

    # df所有的节点是柱线底背离
    def detect_all_macd_column_down(self, span):
        hist = self.df['Histogram'].to_numpy()
        close = self.df['close'].to_numpy()
        date = self.df['trade_date'].to_numpy()
        res = []
        i = 0
        while i < len(hist):
            if self.detect_macd_column_down(hist, close, i, span):
                res.append(date[i])
                #i += span  # 满足条件后跳过span个元素
            #else:
            i += 1
        return res;

    # 检测输入的那一段是否为波峰顶背离信号
    def detect_macd_top_back(self, macd, ma10):
        """
        当macd线出现下降趋势且ma10为上升趋势时
        """
        #print(macd)
        kOfMacd, _ = shapeOfLine().linear_regression(arr = macd)
        kOfma10, _ = shapeOfLine().linear_regression(arr = ma10)
        return kOfMacd < 0 and kOfMacd > 0

    # df所有的点是macd波峰顶背离
    def detect_all_macd_top_back(self, span):
        MACD = self.df['MACD'].to_numpy()
        ma10 = self.df['MA10'].to_numpy()
        date = self.df['trade_date'].to_numpy()
        res = []
        i = span
        while i < len(MACD):
            if self.detect_macd_top_back(macd = MACD[i - span: i], ma10 = ma10[i - span: i]):
                res.append(date[i])
                i += span
            else: i += 1
        return res

    # 检测某个点是否为高位死叉背离
    def detect_macd_top_cross(self, macd, close, dea, target, span):
        """
        当一个点的close运行到顶部，且在周围出现死叉。可以判断为高位死叉
        """
        #print("11")
        if target - span < 0 or target + span > len(macd): return False
        for i in range(target - span, target + span, 1):
            if shapeOfLine.isTop(arr = close, target = i, span = span) : flag = True
        flag = shapeOfLine.isDeadCross(MACD = macd, DEA = dea, target = target)
        return flag

    #df所有的点是macd高位死叉背离
    def detect_all_macd_top_cross(self, span = 5):
        MACD = self.df['MACD'].to_numpy()
        close = self.df['close'].to_numpy()
        dea = self.df['DEA'].to_numpy()
        date = self.df['trade_date'].to_numpy()
        res = []
        #print(len(MACD))
        for i in range(len(MACD)):
            if self.detect_macd_top_cross(macd = MACD, close = close, dea = dea, target = i, span = span) :
                res.append(date[i])
        return res

    def detect_macd_top_column(self, hist, close, target, spanOfclose = 5, spanOfhist = 3):
        """
        当close不断上升且hist的红柱不断下降时
        """
        if target - spanOfclose < 0 or target + spanOfclose >= len(close) or target - spanOfhist < 0 or target + spanOfhist >= len(hist): return False

        flag = True
        for i in range(target - spanOfhist, target, 1):
            if hist[i] <= 0: return False
        k, _ = shapeOfLine().linear_regression(close[target - spanOfclose : target])
        #print(k)
        flag &= (k > 0)
        flag &= shapeOfLine.isMonotony_inpast(arr = hist, ascent = -1, span = spanOfhist, target = target)
        return flag

    def detect_all_macd_top_column(self, spanOfclose = 5, spanOfhist = 3):
        close = self.df['close'].to_numpy()
        hist = self.df['Histogram'].to_numpy()
        date = self.df['trade_date'].to_numpy()
        res = []
        i = 0
        while i < len(hist):
            if self.detect_macd_top_column(hist = hist, close = close, target = i, spanOfclose = spanOfclose, spanOfhist = spanOfhist):
                res.append(date[i])
            i += 1
        return res

    # 画出对应的信号点
    """
    typ
    0  : 波谷底背离
    1  : 死叉底背离
    2  : 柱线底背离
    3  : 波峰顶背离
    4  : 高位死叉背离
    5  : 柱线顶背离
    """
    def plot_all_macd_trough(self, typ):

        plt.rcParams['font.sans-serif'] = ['SimHei']  # 或者['Microsoft YaHei'] 如果你有这个字体
        plt.rcParams['axes.unicode_minus'] = False  # 解决负号无法正常显示的问题
        #df = self.db.get_index_daily(ts_code=ts_code, start=start_date, end=end_date)
        if typ == 0:
            res = self.detect_all_macd_trough_valley()
            typstr = '波谷底背离'
        elif typ == 1:
            res = self.detect_all_macd_cross_down()
            typstr = '死叉底背离'
        elif typ == 2:
            res = self.detect_all_macd_column_down(span = 2)
            typstr = '柱线底背离'
        elif typ == 3:
            res = self.detect_all_macd_top_back(span = 10)
            typstr = '波峰顶背离'
        elif typ == 4:
            res = self.detect_all_macd_top_cross()
            typstr = '高位死叉背离'
        elif typ == 5:
            res = self.detect_all_macd_top_column()
            typstr = '柱线顶背离'
        #print(res)
        fig, (ax1, ax3) = plt.subplots(2, 1, sharex=True, figsize=(24, 16))

        # 绘制MACD和DEA曲线
        ax1.plot(self.df['trade_date'], self.df['MACD'], label='MACD')
        ax1.plot(self.df['trade_date'], self.df['DEA'], label='DEA')
        ax3.plot(self.df['trade_date'], self.df['MA10'], label='MA10')

        #print(len(res))

        #print(type(self.df['trade_date'].iloc[0]))

        for date in res:
            try:
                idx = self.df[self.df['trade_date'] == date].index[0]
                ma10_val = self.df.loc[idx, 'MA10']
                xystart = (date, ma10_val)
                xyend = (date, ma10_val + 0.5)
                ax3.annotate("", xy=xystart, xytext=xyend,
                         arrowprops=dict(arrowstyle='-|>', lw=3, color='black'), annotation_clip=False)

                ax3.text(date, ma10_val + 0.6, str(date) + typstr, rotation=0,
                         va='center', ha='center', color='red', fontsize=8)

            except IndexError:
                pass

        ax1.set_ylabel('MACD & DEA Values')
        ax1.legend()

        newdf = self.df
        newdf['trade_date'] = pd.to_datetime(self.df['trade_date'])  # 确保trade_date列是datetime类型
        newdf['trade_date'] = mdates.date2num(self.df['trade_date'])  # 将datetime转换为Matplotlib datenums
        ohlc = newdf[['trade_date', 'open', 'high', 'low', 'close']].values.tolist()
        candlestick_ohlc(ax3, ohlc, width=0.6, colorup='green', colordown='red')

        # 绘制成交量图
        ax1.bar(self.df['trade_date'], self.df['Histogram'], color='grey', alpha=0.7)
        ax1.set_ylabel('Volume')
        ax1.xaxis.set_major_formatter(DateFormatter('%Y-%m-%d'))  # 设置x轴日期格式

        # 设置图形整体的日期格式和标题
        plt.xticks(rotation=45)  # 旋转x轴标签，避免重叠
        plt.tight_layout()  # 自动调整子图参数，使之填充整个图像区域
        plt.show()

